import cv2,joblib,time,zipfile,os,tempfile
import numpy as np
import util

train_time:float
test_time:float

class PCAKNN:
    def __init__(self):
        self.knn=None
        self.mean_pca=None
        self.eigenvectors=None
        self.k=3
        self.num_components= 50
    
    def train(self,data:np.ndarray,label:np.ndarray):
        """训练PCA-KNN模型
        Args:
            data: 训练数据矩阵，形状为[n_samples, n_features]
            label: 训练标签，形状为[n_samples,]
        """
        self.mean_pca, self.eigenvectors = cv2.PCACompute(data, mean=None, maxComponents=self.num_components)
        data= cv2.PCAProject(data, self.mean_pca, self.eigenvectors)
        self.knn=self.trainKNN(data,label)

    def predict(self,data:np.ndarray):
        """预测数据类别
        Args:
            data: 测试数据矩阵，形状为[n_samples, n_features]
        Returns:
            预测结果数组，形状为[n_samples,]
        """
        data=cv2.PCAProject(data, self.mean_pca, self.eigenvectors)
        _,results,_,_=self.knn.findNearest(data,self.k)
        return results.astype(np.int64).flatten()

    def trainKNN(self,data:np.ndarray,label:np.ndarray):
        knn = cv2.ml.KNearest_create()
        knn.setDefaultK(self.k)
        knn.train(data, cv2.ml.ROW_SAMPLE, label)
        return knn

def launch_train():
    data,label=util.load_vector("number_data/training/")
    pca_knn=PCAKNN()
    global train_time
    start_time=time.time()
    pca_knn.train(data,label)
    train_time=time.time()-start_time
    return pca_knn

def launch_test(model:PCAKNN):
    test_data,test_label=util.load_vector("number_data/testing/")
    global test_time
    start_time=time.time()
    predict_label=model.predict(test_data)
    test_time=time.time()-start_time
    util.show_result(test_label,predict_label)


def save(model: PCAKNN, path: str):
    """将模型保存为单个zip文件（内部包含KNN模型和参数）"""
    with zipfile.ZipFile(path, 'w') as zipf:
        # 将KNN模型写入zip内的指定条目
        with tempfile.TemporaryDirectory() as tmp_dir:
            knn_path = os.path.join(tmp_dir, "knn_model.xml")
            model.knn.save(knn_path)
            zipf.write(knn_path, arcname="knn_model.xml")
        
        # 将其他参数写入zip内的joblib条目
        params = {
            "mean_pca": model.mean_pca,
            "eigenvectors": model.eigenvectors,
            "k": model.k,
            "num_components": model.num_components
        }
        with tempfile.NamedTemporaryFile() as tmp_file:
            joblib.dump(params, tmp_file.name)
            zipf.write(tmp_file.name, arcname="params.joblib")

def load(path: str) -> PCAKNN:
    """从zip文件加载模型"""
    model = PCAKNN()
    with tempfile.TemporaryDirectory() as tmp_dir:
        with zipfile.ZipFile(path, 'r') as zipf:
            # 提取KNN模型文件
            knn_xml_path = os.path.join(tmp_dir, "knn_model.xml")
            zipf.extract("knn_model.xml", path=tmp_dir)
            model.knn = cv2.ml.KNearest_load(knn_xml_path)
            
            # 加载参数文件
            params_path = os.path.join(tmp_dir, "params.joblib")
            zipf.extract("params.joblib", path=tmp_dir)
            params = joblib.load(params_path)
            
        # 填充参数
        model.mean_pca = params["mean_pca"]
        model.eigenvectors = params["eigenvectors"]
        model.k = params["k"]
        model.num_components = params["num_components"]
    return model

def report_time():
    print(f"task1 trainning cost {train_time}s")
    print(f"task1 testing cost {test_time}s")

def main():
    model=launch_train()
    launch_test(model)
    save(model,"result/task1/task1.model")
    report_time()


if __name__=="__main__":
    main()